package com.example.tls_mitm_learn;

import android.net.VpnService;

import org.pcap4j.packet.DnsPacket;
import org.pcap4j.packet.IpPacket;
import org.pcap4j.packet.UdpPacket;
import org.pcap4j.packet.namednumber.DnsResourceRecordType;

import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.net.InetSocketAddress;

import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

public class UdpConnect implements PacketTool.NetworkChannel, BytesParse.Parse {

    private final BytesParse bytesParse;
    private final DatagramSocket datagramSocket;
    private final IpPacket dstIpPacket;
    private final BlockingQueue<IpPacket> queue = new LinkedBlockingQueue<>();

    private static final ExecutorService executor = Executors.newCachedThreadPool();;
    private  Future<?>  readFuture = null;
    private  Future<?>  writeFuture = null;

    private VpnService vpnService;

    public UdpConnect(VpnService vpnService, BytesParse bytesParse , IpPacket ipPacket) throws Exception {
        this.bytesParse = bytesParse;

        dstIpPacket = ipPacket;
        datagramSocket = new DatagramSocket();
        boolean protect = vpnService.protect(datagramSocket);
        datagramSocket.setSoTimeout(3000);
        this.vpnService = vpnService;
        readFuture = executor.submit(this::read);
        writeFuture = executor.submit(this::write);
    }

    @Override
    public void parseIpPacket(IpPacket ipPacket) {
        queue.offer(ipPacket);
    }

    private void read(){
        try{

            UdpPacket udpPacket = (UdpPacket) dstIpPacket.getPayload();
            int dstPort = udpPacket.getHeader().getDstPort().valueAsInt();

            while (!Thread.interrupted()){

                byte[] buffer = new byte[1024];
                DatagramPacket packet = new DatagramPacket(buffer, buffer.length);
                datagramSocket.receive(packet);

                if(dstPort == 53){
                    try {
                        DnsPacket dnsPacket = DnsPacket.newPacket(buffer, 0, packet.getLength());

                        AtomicReference<String> domain = new AtomicReference<>("");

                        dnsPacket.getHeader().getQuestions().forEach(dnsQuestion -> {
                            domain.set(dnsQuestion.getQName().getName());

                        });

                        // 处理响应包，提取 A/AAAA 记录
                        dnsPacket.getHeader().getAnswers().forEach(answer -> {

                            if(answer.getDataType() == DnsResourceRecordType.A){
                                byte[] rawData = answer.getRData().getRawData();
                                if(rawData.length == 4){
                                    String ip = (rawData[0] & 0xFF) + "." +
                                            (rawData[1] & 0xFF) + "." +
                                            (rawData[2] & 0xFF) + "." +
                                            (rawData[3] & 0xFF);
                                    System.out.println("解析到 IP: " + ip);

                                    MyVpnService myVpnService = (MyVpnService) vpnService;
                                    myVpnService.ipAndHostMap.put(ip, domain.get());
                                    System.out.println("[DNS缓存] " + ip + " -> " + domain);
                                }

                            }

                        });

                    } catch (Exception e) {
                        System.err.println("解析 DNS 出错: " + e.getMessage());
                    }
                }

                //Push
                PacketTool.sendUdpPacket(this,buffer);

            }
        }catch (Exception e){
            close();
        }

    }

    private void write(){

        try {

            InetAddress dstAdder = dstIpPacket.getHeader().getDstAddr();
            UdpPacket udpPacket = (UdpPacket) dstIpPacket.getPayload();
            int dstPort = udpPacket.getHeader().getDstPort().valueAsInt();


            while (!Thread.interrupted()){
                IpPacket packet = queue.take();

                UdpPacket uPacket = (UdpPacket) packet.getPayload();
                byte[] rawData = uPacket.getPayload().getRawData();

                DatagramPacket datagramPacket = new DatagramPacket(rawData, rawData.length, dstAdder, dstPort);

                datagramSocket.send(datagramPacket);
            }
        }catch (Exception e){
            close();

        }

    }

    private void close(){

        if(datagramSocket!=null){
            datagramSocket.close();
        }

        if(readFuture!=null){
            //发送中断指令
            readFuture.cancel(true);
        }

        if(writeFuture!=null){
            //发送中断指令
            writeFuture.cancel(true);
        }

        if(dstIpPacket!=null){
            bytesParse.removeKey(dstIpPacket);
        }

        if (dstIpPacket != null) {
            bytesParse.removeKey(dstIpPacket);
        }


    }

    @Override
    public AtomicInteger getSequenceNumber() {
        return null;
    }

    @Override
    public AtomicInteger getAckNumber() {
        return null;
    }

    @Override
    public void sendToVpn(byte[] packetData) {
        bytesParse.pushData(packetData);
    }

    @Override
    public InetSocketAddress getSrcAddress() {
        InetAddress dstAdder = dstIpPacket.getHeader().getDstAddr();
        UdpPacket udpPacket = (UdpPacket) dstIpPacket.getPayload();
        int dstPort = udpPacket.getHeader().getDstPort().valueAsInt();

        return new InetSocketAddress(dstAdder, dstPort);
    }

    @Override
    public InetSocketAddress getDstAddress() {
        InetAddress srcAdder = dstIpPacket.getHeader().getSrcAddr();
        UdpPacket udpPacket = (UdpPacket) dstIpPacket.getPayload();
        int srcPort = udpPacket.getHeader().getSrcPort().valueAsInt();
        return new InetSocketAddress(srcAdder, srcPort);
    }


}
